#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python常用模块库 
@File ：file_generator.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/2/21 上午9:35 
"""

"""
测试文件生成器
Version: 1.0
功能：
- 生成指定数量的文件
- 生成指定大小的文件
- 文件名包含精确时间戳
- 支持并行生成
- 生成进度显示
"""

import argparse
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path


class TestFileGenerator:
	"""测试文件生成器"""

	def __init__(self,
				 output_dir: str = "./test_files",
				 file_prefix: str = "test",
				 max_workers: int = 4):
		"""
		初始化生成器
		:param output_dir: 输出目录
		:param file_prefix: 文件名前缀
		:param max_workers: 最大并发数
		"""
		self.output_dir = Path(output_dir).resolve()
		self.file_prefix = file_prefix
		self.max_workers = max_workers
		self._setup_logging()
		self._prepare_directory()

	def _setup_logging(self):
		"""配置日志系统"""
		logging.basicConfig(
			level=logging.INFO,
			format='%(asctime)s [%(levelname)s] %(message)s',
			datefmt='%Y-%m-%d %H:%M:%S',
			handlers=[
				logging.FileHandler('file_generator.log'),
				logging.StreamHandler()
			]
		)
		self.logger = logging.getLogger(__name__)

	def _prepare_directory(self):
		"""准备输出目录"""
		try:
			self.output_dir.mkdir(parents=True, exist_ok=True)
			self.logger.info(f"输出目录已准备: {self.output_dir}")
		except Exception as e:
			self.logger.error(f"创建目录失败: {str(e)}")
			raise

	def generate_files(self,
					   file_count: int = 100,
					   file_size_mb: float = 1.0,
					   time_interval: float = 0.1):
		"""
		生成测试文件
		:param file_count: 文件数量
		:param file_size_mb: 单个文件大小(MB)
		:param time_interval: 时间间隔(秒)
		"""
		start_time = time.time()
		self.logger.info(f"开始生成 {file_count} 个文件 | 单文件大小: {file_size_mb}MB")

		# 计算总数据量
		total_size_gb = (file_count * file_size_mb) / 1024
		self.logger.info(f"预计总大小: {total_size_gb:.2f}GB")

		# 使用线程池并行生成
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
			futures = []
			for i in range(file_count):
				futures.append(
					executor.submit(
						self._generate_single_file,
						file_index=i + 1,
						file_size_mb=file_size_mb,
						time_interval=time_interval
					)
				)

			# 进度跟踪
			completed = 0
			for future in futures:
				future.result()
				completed += 1
				if completed % 10 == 0:
					self.logger.info(f"生成进度: {completed}/{file_count} ({completed / file_count:.1%})")

		elapsed = time.time() - start_time
		self.logger.info(f"文件生成完成! 总耗时: {elapsed:.2f}秒 | 平均速度: {total_size_gb / elapsed:.2f}GB/s")

	def _generate_single_file(self,
							  file_index: int,
							  file_size_mb: float,
							  time_interval: float):
		"""生成单个文件"""
		try:
			# 生成时间戳文件名
			timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
			filename = f"{self.file_prefix}_{timestamp}_{file_index:06d}.dat"
			file_path = self.output_dir / filename

			# 生成文件内容
			file_size = int(file_size_mb * 1024 * 1024)
			self._write_file(file_path, file_size)

			# 模拟时间间隔
			time.sleep(time_interval)
			return file_path
		except Exception as e:
			self.logger.error(f"生成文件失败: {str(e)}")
			raise

	def _write_file(self, path: Path, size: int):
		"""写入指定大小的文件"""
		# 使用稀疏文件提高大文件生成效率
		with open(path, 'wb') as f:
			f.seek(size - 1)
			f.write(b'\0')
		self.logger.debug(f"已生成文件: {path.name} ({size / 1024 / 1024:.2f}MB)")


def parse_args():
	"""解析命令行参数"""
	parser = argparse.ArgumentParser(description="测试文件生成器")
	parser.add_argument("-n", "--num", type=int, default=100,
						help="生成文件数量 (默认: 100)")
	parser.add_argument("-s", "--size", type=float, default=1.0,
						help="单个文件大小(MB) (默认: 1.0)")
	parser.add_argument("-o", "--output", type=str, default="./test_files",
						help="输出目录 (默认: ./test_files)")
	parser.add_argument("-w", "--workers", type=int, default=4,
						help="并发工作数 (默认: 4)")
	parser.add_argument("-i", "--interval", type=float, default=0.1,
						help="生成时间间隔(秒) (默认: 0.1)")
	return parser.parse_args()


if __name__ == "__main__":
	try:
		args = parse_args()

		generator = TestFileGenerator(
			output_dir=args.output,
			max_workers=args.workers
		)

		generator.generate_files(
			file_count=args.num,
			file_size_mb=args.size,
			time_interval=args.interval
		)
	except Exception as e:
		logging.error(f"文件生成失败: {str(e)}", exc_info=True)
